import numpy as np
import os
import tensorflow as tf
import cv2
from data.config import cfg
from utils import utils

print(cfg.train.train_path)
class Dataset(object):
    """
        实现训练前的数据处理
    """
    def __init__(self, datatype):
        self.annotation_path = cfg.train.train_path if datatype == 'train' else cfg.test.test_path
        self.batch_size = cfg.train.batch_size
        self.input_size = cfg.train.input_size
        # yolo中的三种anchor比例
        self.strides = np.array(cfg.yolo.strides)
        self.anchors = utils.get_anchors(cfg.yolo.anchors_path)
        self.num_per_anchor_scale = self.anchors.shape[1] # 同尺寸下的不同比例的个数
        self.maxAnchorNum_per_scale = cfg.yolo.maxAnchorNum_per_scale  # 生成标签框时，每一种比例下最多生成的锚框
        # 获取标签名称
        self.classes = utils.get_class_names(cfg.yolo.class_path)
        self.num_class = len(self.classes)

        # 加载数据集
        self.annotations = self.load_annotations()
        self.example_num = len(self.annotations)
        self.batchs_num = int(np.ceil(self.example_num / self.batch_size))
        self.batch_count = 0  # 默认第0个批次


    def load_annotations(self):
        assert os.path.exists(self.annotation_path),'cannot find {}'.format(self.annotation_path)
        with open(self.annotation_path,'r') as f:
            lines = f.readlines()
            annotations = [line.strip() for line in lines if len(line.strip().split()[1:]) !=0]
        np.random.shuffle(annotations)
        print(annotations)
        return annotations



    def __len__(self):
        pass

    def __iter__(self):
        return self

    def __next__(self):
        with tf.device('/cpu:0'):
            # 根据输入图片的大小，计算模型的输出的图片大小，若输入为416*416,则会输出
            self.train_output_size = self.input_size // self.strides

            # 构建输出图片跟数据的shape 如：(n,416,416,3)大小为416的三通道图
            batch_images = np.zeros((self.batch_size, self.input_size, self.input_size, 3))

            # 分别构建批量情况下，三种不同尺寸下 真实框对应的标签信息label_xxx，记忆真实框的信息xxx
            batch_label_sbboxes = np.zeros((self.batch_size,
                                            self.train_output_size[0],
                                            self.train_output_size[0],
                                            self.num_per_anchor_scale,
                                            5 + self.num_class), dtype=np.float32)
            batch_label_mbboxes = np.zeros((self.batch_size,
                                            self.train_output_size[1],
                                            self.train_output_size[1],
                                            self.num_per_anchor_scale,
                                            5 + self.num_class), dtype=np.float32)
            batch_label_lbboxes = np.zeros((self.batch_size,
                                            self.train_output_size[2],
                                            self.train_output_size[2],
                                            self.num_per_anchor_scale,
                                            5 + self.num_class), dtype=np.float32)
            batch_sbboxes = np.zeros((self.batch_size, self.maxAnchorNum_per_scale, 4), dtype=np.float32)
            batch_mbboxes = np.zeros((self.batch_size, self.maxAnchorNum_per_scale, 4), dtype=np.float32)
            batch_lbboxes = np.zeros((self.batch_size, self.maxAnchorNum_per_scale, 4), dtype=np.float32)


            num = 0
            if self.batch_count < self.batchs_num:
                while num < self.batch_size:
                    index = self.batch_count * self.batch_size + num
                    # 下标超出样本个数
                    if index > self.example_num: index -= self.example_num
                    # 每一张图片信息
                    annotation = self.annotations[index]

                    # 每一张图片的缩放, 裁剪等基础操作 缩放到根输入要求大小一致
                    image,bboxes = self.parse_annotation(annotation)


                    # 根据真实框，生成对应的标签
                    label_sbbox, label_mbbox, label_lbbox, sbbox, mbbox, lbbox = self.preprocess_true_boxes(bboxes)

                    # 批量分装
                    batch_images[num, :, :, :] = image
                    batch_label_sbboxes[num, :, :, :, :] = label_sbbox
                    batch_label_mbboxes[num, :, :, :, :] = label_mbbox
                    batch_label_lbboxes[num, :, :, :, :] = label_lbbox
                    batch_sbboxes[num, :, :] = sbbox
                    batch_mbboxes[num, :, :] = mbbox
                    batch_lbboxes[num, :, :] = lbbox

                    num += 1
                self.batch_count += 1

                batch_target_smaller = batch_label_sbboxes, batch_sbboxes
                batch_target_middle = batch_label_mbboxes, batch_mbboxes
                batch_target_larger = batch_label_lbboxes, batch_lbboxes

                return (batch_images,
                        (batch_target_smaller, batch_target_middle, batch_target_larger))

            else:
                self.batch_count = 0
                np.random.shuffle(self.annotations)
                raise StopIteration

    # 解析每一张图片信息，包括缩放到指定输入大小，等等操作
    def parse_annotation(self, annotation):
        """
            解析每一张图片信息，对图片及box 做一系列操作
        :param annotation:
        :return: 处理后的图片以及 图片上框的图片
        """
        line = annotation.split()
        image_path = line[0]
        assert os.path.exists(image_path),'the {} is not exist'.format(image_path)

        image = cv2.imread(image_path)
        bboxes = [list(map(lambda x: int(float(x)), box.split(','))) for box in line[1:]]

        # 此处可以进行颜色翻转、水平翻转，裁剪，放射变换等操作


        # 图片的缩放操作, 缩放到指定输入图片的大小
        image, bboxes = utils.image_preprocess(image, (self.input_size, self.input_size), bboxes)


        # 此处可以对坐标做合理性判断
        #  省略

        return image, bboxes

    # 针对所有真实框，进行标签处理
    def preprocess_true_boxes(self, gt_box):
        """
            遍历每个真实框，并为每个真实框生成9个anchor框，进行比较，标记正负样本，
        :param gt_box:
        :return:分别返回三种不同尺寸下，各像素点上对应的值 如 label_sbbox(52,52,3,5+self.num_class),其中第三个维度代表三个不同比例
                sbbox(150,4)为每个尺寸下，真实框的坐标值，当前尺寸下，最多有150个框；其它尺寸的同理
        """

        # 三种不同尺寸下 三个不同比例anchor框的shape
        label = [
            np.zeros(
                (
                    self.train_output_size[i],
                    self.train_output_size[i],
                    self.num_per_anchor_scale,
                    5 + self.num_class,
                )
            )
            for i in range(3)
        ]
        bboxes_xywh = [np.zeros((self.maxAnchorNum_per_scale, 4)) for _ in range(3)]
        bbox_count = np.zeros((3,))  # 记录三种不同尺寸下，各存在正样本的个数

        for bbox in gt_box:
            # 获取真实框的左上，右下顶点坐标 以及框中物体的类别
            bbox_coor = bbox[0:4]
            bbox_class_index = bbox[4]

            # 将物体类别进行热编码 并进行label_smoothing操作
            onehot = np.zeros(self.num_class, dtype=np.float)
            onehot[bbox_class_index] = 1.0
            uniform_distribution = np.full(self.num_class, 1.0/self.num_class)
            deta = 0.01
            smooth_onehot = onehot * (1 - deta) + uniform_distribution * deta




            # 计算真实框的中心点和宽高 (4,)
            bbox_xywh = np.concatenate([(bbox_coor[0:2] + bbox_coor[2:]) * 0.5, (bbox_coor[2:0] - bbox_coor[0:2])], axis=-1)

            # 按照三种不同大小的尺寸缩放真实框(1,4)/(3,1)--->(3,4)
            bbox_xywh_scaled = 1.0 * bbox_xywh[np.newaxis, :] / self.strides[:, np.newaxis]

            # 针对同一个真实框，以真实框的中心点为中心，根据3种不同尺寸，3种不同比构建9个anchor框
            iou = []
            exist_positive = False  # 默认真实框与anchor框的iou都小于0.3，即默认把当前框标记为负样本
            for i in range(len(self.strides)):
                # 在当前的尺寸大小下，构建三种不同比例(1:1,1:2,2:1)的anchor框坐标
                anchors_xywh = np.zeros((self.num_per_anchor_scale, 4))

                # 构建三种不同比例尺(1:1,1:2,2:1)下,anchor框的中心点 就是真实框的中心点
                anchors_xywh[:, 0:2] = np.float(bbox_xywh_scaled[i][0:2]).astype(np.int32) + 0.5

                # 三种不同比例尺(1:1,1:2,2:1)下，anchor框的长宽 self.anchors[i]为三种不同比例下的框的长宽
                anchors_xywh[:, 2:4] = self.anchors[i]

                # 当前尺寸下, 计算三个anchor框与真实框的iou 此处(1,4)和(3,4)--->(3,)
                iou_scale = utils.bbox_iou(bbox_xywh_scaled[i][np.newaxis, :], anchors_xywh)

                iou.append(iou_scale)
                iou_mask = iou_scale > 0.3

                # 当前尺寸下，三个比例的anchor框与真实框的iou有一个是符合条件的
                if np.any(iou_mask):
                    # 按比例缩放后真实框中心所在像素网格的x,y坐标
                    xgrid, ygrid = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32)

                    # 分别记录三种尺寸下，真实框坐在像素位置，最大iou位置对应真实框信息
                    label[i][ygrid, xgrid, iou_mask, :] = 0
                    label[i][ygrid, xgrid, iou_mask, 0:4] = bbox_xywh
                    label[i][ygrid, xgrid, iou_mask, 4:5] = 1.0
                    label[i][ygrid, xgrid, iou_mask, 5] = smooth_onehot

                    # 分别记录三种尺寸下，所有的真实框
                    index = int(bbox_count[i] % self.maxAnchorNum_per_scale)
                    bboxes_xywh[i][index, 0:4] = bbox_xywh
                    bbox_count[i] += 1

                    exist_positive = True


            # 当前框为负样本时
            if not exist_positive:
                # 找到iou中最大值对应的下标 与真实框iou最大的anchor对应的下标，此处应先把anchor对应的iou按尺寸顺序展平成一行
                best_anchor_index = np.argmax(np.array(iou).reshape(-1), axis=-1)
                # 求此anchor对应的的是哪个尺寸（[52,26,13]）
                best_anchor_size_index = int(best_anchor_index / self.num_per_anchor_scale)
                # 求此anchor对应的的是哪个比例(1: 1, 1: 2, 2: 1)
                best_anchor_ratio_index = int(best_anchor_index % self.num_per_anchor_scale)

                # 计算此anchor对应尺寸下， 真实框中心所在像素网格的x,y坐标
                xgrid, ygrid = np.floor(bbox_xywh_scaled[best_anchor_size_index, 0:2]).astype(np.int32)

                # 分别记录三种尺寸下，真实框坐在像素位置，最大iou位置对应真实框信息
                label[best_anchor_size_index][ygrid, xgrid, best_anchor_ratio_index, :] = 0
                label[best_anchor_size_index][ygrid, xgrid, best_anchor_ratio_index, 0:4] = bbox_xywh
                label[best_anchor_size_index][ygrid, xgrid, best_anchor_ratio_index, 4:5] = 1.0
                label[best_anchor_size_index][ygrid, xgrid, best_anchor_ratio_index, 5] = smooth_onehot

                # 分别记录三种尺寸下，所有的真实框
                index = int(bbox_count[best_anchor_size_index] % self.maxAnchorNum_per_scale)
                bboxes_xywh[best_anchor_size_index][index, 0:4] = bbox_xywh
                bbox_count[best_anchor_size_index] += 1
        label_sbbox, label_mbbox, label_lbbox = label
        sbbox, mbbox, lbbox = bboxes_xywh
        return label_sbbox, label_mbbox, label_lbbox, sbbox, mbbox, lbbox









if __name__ == '__main__':
    # dataset = Dataset('train')
    # print(dataset)
    pass
